use std::array;

use super::{br, lfsr::Lfsr, nl_func::NonlinearFunction};

/// 使用初始密钥k和初始化向量iv,
/// 获取一个可以不断生成ZUC密钥的闭包.
pub fn get_zuc_key_generator(k: [u8; 16], iv: [u8; 16]) -> impl FnMut() -> u32 {
    let mut zuc = Zuc::with_key(k, iv);
    zuc.init();
    zuc.work_prepare();
    move || zuc.get_next_key()
}

/// ZUC算法密钥流生成器.
#[derive(Debug)]
struct Zuc {
    lfsr: Lfsr,
    nonlinear_func: NonlinearFunction,
}

impl Zuc {
    /// 使用初始密钥k和初始化向量iv构建ZUC.
    ///
    /// 返回的对象还需要调用`init`和`work_prepare`才可以进行密钥输出.
    pub fn with_key(k: [u8; 16], iv: [u8; 16]) -> Self {
        let s = array::from_fn(|i| ((k[i] as u32) << 23) + (D[i] << 8) + iv[i] as u32);
        Self {
            lfsr: Lfsr::with_array(s),
            nonlinear_func: NonlinearFunction::new(),
        }
    }

    /// 初始化阶段.
    pub fn init(&mut self) {
        for _i in 0..32 {
            let (x0, x1, x2, _) = self.bit_reconstruct();
            let w = self.f(x0, x1, x2);
            let u = w >> 1;
            self.lfsr.lfsr_with_initialization_mode(u);
            // println!("DEBUG: - {_i}");
            // println!(
            //     "DEBUG: x0: {:x?}, x1: {:x}, x2: {:x}, w: {:x}",
            //     x0, x1, x2, w
            // );
            // println!("DEBUG: nf: {:?}", self.nonlinear_func);
            // println!("DEBUG: S15: {:x}", self.lfsr.get_s().get(15));
        }
    }

    /// 工作阶段的第一步.
    /// 执行完毕后应当已可以输出密钥.
    pub fn work_prepare(&mut self) {
        let (x0, x1, x2, _) = self.bit_reconstruct();
        let _ = self.f(x0, x1, x2);
        self.lfsr.lfsr_with_work_mode();
    }

    /// 生成一个 32bit 密钥.
    pub fn get_next_key(&mut self) -> u32 {
        let (x0, x1, x2, x3) = self.bit_reconstruct();
        let z = self.f(x0, x1, x2) ^ x3;
        self.lfsr.lfsr_with_work_mode();
        z
    }

    /// 从lfsr内的寄存器组s进行比特重组.
    ///
    /// 输出 `(x0, x1, x2, x3)`.
    fn bit_reconstruct(&self) -> (u32, u32, u32, u32) {
        br::bit_reconstruct(self.lfsr.get_s())
    }

    /// 非线性函数加工.
    fn f(&mut self, x0: u32, x1: u32, x2: u32) -> u32 {
        self.nonlinear_func.f(x0, x1, x2)
    }
}

const D: [u32; 16] = [
    0b100010011010111,
    0b010011010111100,
    0b110001001101011,
    0b001001101011110,
    0b101011110001001,
    0b011010111100010,
    0b111000100110101,
    0b000100110101111,
    0b100110101111000,
    0b010111100010011,
    0b110101111000100,
    0b001101011110001,
    0b101111000100110,
    0b011110001001101,
    0b111100010011010,
    0b100011110101100,
];

mod tests {

    #[test]
    fn test_zuc() {
        use super::Zuc;

        let k = [
            0x3d, 0x4c, 0x4b, 0xe9, 0x6a, 0x82, 0xfd, 0xae, 0xb5, 0x8f, 0x64, 0x1d, 0xb1, 0x7b,
            0x45, 0x5b,
        ];
        let iv = [
            0x84, 0x31, 0x9a, 0xa8, 0xde, 0x69, 0x15, 0xca, 0x1f, 0x6b, 0xda, 0x6b, 0xfb, 0xd8,
            0xc7, 0x66,
        ];

        let mut zuc = Zuc::with_key(k, iv);
        // 检查最初的s状态.
        let target_s = [
            0x1ec4d784, 0x2626bc31, 0x25e26b9a, 0x74935ea8, 0x355789de, 0x4135e269, 0x7ef13515,
            0x5709afca, 0x5acd781f, 0x47af136b, 0x326bc4da, 0x0e9af16b, 0x58de26fb, 0x3dbc4dd8,
            0x22f89ac7, 0x2dc7ac66,
        ];
        let s = zuc.lfsr.get_s();
        for i in 0..16 {
            println!("origin s[{i}]: {:x}", s.get(i));
            assert_eq!(*s.get(i), target_s[i]);
        }

        zuc.init();
        // 检查初始化后s状态
        let target_s = [
            0x10da5941, 0x5b6acbf6, 0x17060ce1, 0x35368174, 0x5cf4385a, 0x479943df, 0x2753bab2,
            0x73775d6a, 0x43930a37, 0x77b4af31, 0x15b2e89f, 0x24ff6e20, 0x740c40b9, 0x026a5503,
            0x194b2a57, 0x7a9a1cff,
        ];
        let s = zuc.lfsr.get_s();
        for i in 0..16 {
            println!("inited s[{i}]: {:x}", s.get(i));
            assert_eq!(*s.get(i), target_s[i]);
        }

        zuc.work_prepare();

        let z1 = zuc.get_next_key();
        println!("z1: {:x}", z1);
        assert_eq!(z1, 0x14f1c272);

        let z2 = zuc.get_next_key();
        println!("z2: {:x}", z2);
        assert_eq!(z2, 0x3279c419);
    }

    #[test]
    fn test_zuc_key_generator_all_0() {
        use super::get_zuc_key_generator;
        let k = [0u8; 16];
        let iv = [0u8; 16];
        let mut zuc_key_generator = get_zuc_key_generator(k, iv);

        assert_eq!(zuc_key_generator(), 0x27bede74);
        assert_eq!(zuc_key_generator(), 0x018082da);
    }

    #[test]
    fn test_zuc_key_generator_all_1() {
        use super::get_zuc_key_generator;
        let k = [0xffu8; 16];
        let iv = [0xffu8; 16];
        let mut zuc_key_generator = get_zuc_key_generator(k, iv);

        assert_eq!(zuc_key_generator(), 0x0657cfa0);
        assert_eq!(zuc_key_generator(), 0x7096398b);
    }
}
